#cnn模型训练代码，训练的代码会保存在models目录下，折线图会保存在results目录下
import tensorflow as tf
import matplotlib.pyplot as plt
from time import *

# 数据集加载函数，指明数据集的位置并统一处理为imgheight*imgwidth的大小，同时设置batch一批处理的个数
def data_load(data_dir, test_data_dir, img_height, img_width, batch_size):
    # 加载训练集
    train_ds = tf.keras.preprocessing.image_dataset_from_directory(
        data_dir, # 指定训练集数据路径
        label_mode='categorical', # 进行独热编码
        seed=123, # 保证随机生成的随机数前后一致。
        image_size=(img_height, img_width),# 对图像risize进行预处理
        batch_size=batch_size) # 每次迭代取16个数据
    # 加载测试集
    val_ds = tf.keras.preprocessing.image_dataset_from_directory(
        test_data_dir,
        label_mode='categorical',
        seed=123,
        image_size=(img_height, img_width),
        batch_size=batch_size)
    class_names = train_ds.class_names
    # 返回处理之后的训练集、验证集和类名
    return train_ds, val_ds, class_names


# 构建CNN模型，class_num设置17种分类
def model_load(IMG_SHAPE=(224, 224, 3), class_num=17):
    # 搭建模型
    model = tf.keras.models.Sequential([
        # 对模型做归一化的处理，将像数值0-255之间的数字统一处理到0到1之间，方便模型后面进行调整
        tf.keras.layers.experimental.preprocessing.Rescaling(1. / 255, input_shape=IMG_SHAPE),
        # 卷积层，该卷积层的输出为32个通道，卷积核的大小是3*3，激活函数为relu，
        # relu函数看起来和行为都像一个线性函数，但实际上是一个非线性函数，允许学习数据中的复杂关系 。该函数还必须提供更灵敏的激活和输入，避免饱和。
        tf.keras.layers.Conv2D(32, (3, 3), activation='relu'),
        # 添加池化层，池化的kernel大小是2*2
        tf.keras.layers.MaxPooling2D(2, 2),
        # Add another convolution
        # 卷积层，输出为64个通道，卷积核大小为3*3，激活函数为relu
        tf.keras.layers.Conv2D(64, (3, 3), activation='relu'),
        # 池化层，最大池化，对2*2的区域进行池化操作
        tf.keras.layers.MaxPooling2D(2, 2),
        # 将二维的输出转化为一维
        tf.keras.layers.Flatten(),
        #全连接层，与预卷积示例中相同的128个密集层和10个输出层：
        tf.keras.layers.Dense(128, activation='relu'),
        # 通过softmax函数将模型输出为类名长度的神经元上，激活函数采用softmax函数对应概率值：
        # 可以对输出值进行归一化操作，把所有输出值都转化为概率，所有概率值加起来等于1
        tf.keras.layers.Dense(class_num, activation='softmax')
    ])
    # 输出模型信息
    model.summary()
    # 指明模型的训练参数，优化器为sgd优化器随机梯度下降，损失函数为交叉熵损失函数
    model.compile(optimizer='sgd', loss='categorical_crossentropy', metrics=['accuracy'])
    # 返回模型
    return model


# 展示训练过程的曲线
def show_loss_acc(history):
    # 从history中提取模型训练集和验证集准确率信息和误差信息
    acc = history.history['accuracy']
    val_acc = history.history['val_accuracy']
    loss = history.history['loss']
    val_loss = history.history['val_loss']

    # 按照上下结构将图画输出
    plt.figure(figsize=(8, 8))#指定图像的宽和高
    plt.subplot(2, 1, 1)#指定图像位置
    plt.plot(acc, label='Training Accuracy')#指定图像x，y轴
    plt.plot(val_acc, label='Validation Accuracy')
    plt.legend(loc='lower right')#图像右上角的标注
    plt.ylabel('Accuracy')#图像y轴名称
    plt.ylim([min(plt.ylim()), 1])#设置x轴/y轴的数值显示范围。
    plt.title('Training and Validation Accuracy')#图像标题

    plt.subplot(2, 1, 2)
    plt.plot(loss, label='Training Loss')
    plt.plot(val_loss, label='Validation Loss')
    plt.legend(loc='upper right')
    plt.ylabel('Cross Entropy')
    plt.title('Training and Validation Loss')
    plt.xlabel('epoch')
    plt.savefig('results/results_cnn.png', dpi=100)#保存图片路径。


def train(epochs):
    # 开始训练，记录开始时间
    begin_time = time()
    # todo 加载数据集， 修改为你的数据集的路径
    # 将图片统一处理成244*244的大小，送入到模型中训练。然后模型一次训练16张图片，可以防止模型过拟合，还可以训练快一点。
    train_ds, val_ds, class_names = data_load("C:/Users/Administrator/Desktop/bbb/newData/train",
                                              "C:/Users/Administrator/Desktop/bbb/newData/test", 224, 224, 16)
    print(class_names)
    # 加载模型
    model = model_load(class_num=len(class_names))
    # 开始训练模型，指明训练的轮数epoch，开始训练
    history = model.fit(train_ds, validation_data=val_ds, epochs=epochs)
    # todo 保存模型， 修改为你要保存的模型的名称
    model.save("models/cnn_flowers.h5")
    # 记录结束时间
    end_time = time()
    run_time = end_time - begin_time
    print('该循环程序运行时间：', run_time, "s")  # 该循环程序运行时间： 1.4201874732
    # 绘制模型训练过程图
    show_loss_acc(history)


if __name__ == '__main__':
    train(epochs=30)
    
    